from nicegui import ui, app
import asyncio
import datetime
import os
import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from std_msgs.msg import Float32MultiArray, Int32
import threading
import json
import math


# 自定义GaitParam消息类
class GaitParam:
    def __init__(self):
        self.gait_type = 0  # uint8
        self.velocity = 0.0  # float32
        self.nominal_height = 0.25  # float32
        self.stance_duration = 0.2  # float32
        self.stance_depth = 0.0  # float32
        self.swing_height = 0.07  # float32
        self.movement = 0  # uint8

    def to_dict(self):
        return {
            'gait_type': self.gait_type,
            'velocity': self.velocity,
            'nominal_height': self.nominal_height,
            'stance_duration': self.stance_duration,
            'stance_depth': self.stance_depth,
            'swing_height': self.swing_height,
            'movement': self.movement
        }


# ROS2配置
ROS_NODE_NAME = "nicegui_go2_controller"
GAIT_MSG_TOPIC = "/go2_gait_planner/gait_msg"
JOINTS_TOPIC = "/go2_gait_planner/joints"
LOWCMD_TOPIC = "/lowcmd"

# Web配置
API_BASE = "http://localhost:8000"
VIDEO_WS_URL = "ws://localhost:8000/ws/video"
DEFAULT_VIDEO_PATH = r"/home/pc/桌面/New/1.mp4"

# 全局状态
robot_status = {
    "position": {"x": 0, "y": 0, "z": 0},
    "target_position": {"x": 0, "y": 0, "z": 0},  # 新增目标位置用于平滑移动
    "joints": [0.0] * 12,
    "gait_params": GaitParam().to_dict(),
    "lowcmd": {
        "mode_pr": 0,
        "mode_machine": 6,
        "motor_cmd": [{
            "mode": 0,
            "q": 0.0,
            "dq": 0.0,
            "tau": 0.0,
            "kp": 0.0,
            "kd": 0.0,
            "reserve": 0,
            "reserve_array": [0] * 4
        }] * 35,
        "crc": 0
    },
    "movement_state": 0  # 0=停止,1=前进,2=后退,3=左移,4=右移
}

video_ws = None
video_container = None
status_label = None
lowcmd_table = None


class GaitSubscriber(Node):
    def __init__(self):
        super().__init__('gait_subscriber')
        self.subscription = self.create_subscription(
            Float32MultiArray,
            GAIT_MSG_TOPIC,
            self.gait_callback,
            10)
        self.get_logger().info(f"已订阅话题: {GAIT_MSG_TOPIC}")

        self.joints_sub = self.create_subscription(
            Float32MultiArray,
            JOINTS_TOPIC,
            self.joints_callback,
            10)

        self.lowcmd_sub = self.create_subscription(
            Float32MultiArray,
            LOWCMD_TOPIC,
            self.lowcmd_callback,
            10)

    def gait_callback(self, msg):
        if len(msg.data) >= 7:
            gait_params = GaitParam()
            gait_params.gait_type = int(msg.data[0])
            gait_params.velocity = float(msg.data[1])
            gait_params.nominal_height = float(msg.data[2])
            gait_params.stance_duration = float(msg.data[3])
            gait_params.stance_depth = float(msg.data[4])
            gait_params.swing_height = float(msg.data[5])
            gait_params.movement = int(msg.data[6])

            robot_status['gait_params'] = gait_params.to_dict()

    def joints_callback(self, msg):
        robot_status['joints'] = list(msg.data)

    def lowcmd_callback(self, msg):
        if len(msg.data) >= 56:
            robot_status['lowcmd'] = {
                "mode_pr": int(msg.data[0]),
                "mode_machine": int(msg.data[1]),
                "motor_cmd": [{
                    "mode": int(msg.data[2 + i * 8]),
                    "q": float(msg.data[3 + i * 8]),
                    "dq": float(msg.data[4 + i * 8]),
                    "tau": float(msg.data[5 + i * 8]),
                    "kp": float(msg.data[6 + i * 8]),
                    "kd": float(msg.data[7 + i * 8]),
                    "reserve": int(msg.data[8 + i * 8]),
                    "reserve_array": [int(msg.data[9 + i * 8 + j]) for j in range(4)]
                } for i in range(35)],
                "crc": int(msg.data[-1])
            }


def ros_thread():
    try:
        rclpy.init()
        gait_subscriber = GaitSubscriber()
        rclpy.spin(gait_subscriber)
        gait_subscriber.destroy_node()
        rclpy.shutdown()
    except Exception as e:
        print(f"ROS线程错误: {str(e)}")


@ui.page('/')
async def index():
    global status_label, video_container, lowcmd_table

    if 'video_source' not in app.storage.user:
        if os.path.exists(DEFAULT_VIDEO_PATH):
            app.storage.user['video_source'] = DEFAULT_VIDEO_PATH
            app.storage.user['auto_play'] = True
        else:
            app.storage.user['video_source'] = ""
            ui.notify(f"默认视频文件 {DEFAULT_VIDEO_PATH} 不存在", type='warning')

    # 主界面布局 - 使用3行网格布局
    with ui.grid(rows=3).classes("w-full h-screen gap-4 p-4"):
        # 第一行：视频区域 (2列)
        with ui.grid(columns=2).classes("w-full h-[30vh] gap-4"):
            # 左边：推拉流视频
            with ui.card().classes("w-full h-full overflow-hidden"):
                ui.label("实时视频流").classes("text-xl")
                ui.add_body_html('''
                    <script src="https://cdn.jsdelivr.net/npm/flv.js@latest"></script>
                    <style>
                        #flv-container { width:100%; height:100% }
                        #flv-player { width:100%; height:100%; object-fit: contain }
                    </style>
                ''')
                video_container = ui.html('''
                    <div id="flv-container">
                        <video id="flv-player" controls muted></video>
                    </div>
                ''').classes("w-full h-[calc(100%-40px)]")
                ui.run_javascript('''
                    function initPlayer() {
                        const videoElement = document.getElementById('flv-player');
                        if (flvjs.isSupported()) {
                            const flvPlayer = flvjs.createPlayer({
                                type: 'flv',
                                url: 'http://127.0.0.1/live/test.live.flv',
                                isLive: true,
                                stashInitialSize: 128,
                                enableWorker: true,
                                lazyLoadMaxDuration: 3 * 60,
                                seekType: 'range'
                            });
                            flvPlayer.attachMediaElement(videoElement);
                            flvPlayer.load();
                            const playPromise = videoElement.play();
                            if (playPromise !== undefined) {
                                playPromise.catch(e => {
                                    console.log('自动播放被阻止，需要用户交互');
                                });
                            }
                            flvPlayer.on(flvjs.Events.ERROR, (errType, errDetail) => {
                                console.error('播放错误:', errType, errDetail);
                                if(errType === flvjs.ErrorTypes.NETWORK_ERROR) {
                                    setTimeout(() => {
                                        flvPlayer.unload();
                                        flvPlayer.detachMediaElement();
                                        flvPlayer.attachMediaElement(videoElement);
                                        flvPlayer.load();
                                        videoElement.play();
                                    }, 3000);
                                }
                            });
                            window.flvPlayer = flvPlayer;
                        } else {
                            console.error('浏览器不支持FLV播放');
                            videoElement.innerHTML = '您的浏览器不支持视频播放';
                        }
                    }
                    setTimeout(initPlayer, 500);
                ''')

            # 右边：本地视频
            with ui.card().classes("w-full h-full overflow-hidden"):
                ui.label("本地视频").classes("text-xl")
                video_player = ui.video("").classes("w-full h-[calc(100%-40px)]").props("controls autoplay muted")
                if app.storage.user['video_source']:
                    video_player.set_source(DEFAULT_VIDEO_PATH)

        # 第二行：控制区域 (2列)
        with ui.grid(columns=2).classes("w-full h-[20vh] gap-4"):
            # 左边：控制按钮
            with ui.card().classes("w-full h-full overflow-auto"):
                ui.label("机器人控制").classes("text-xl")
                status_label = ui.label(f"状态: 位置(0.00, 0.00)").classes("text-lg")
                with ui.row().classes("w-full gap-2"):
                    ui.button("前进", on_click=lambda: set_movement_state(1)).classes("bg-green-500 hover:bg-green-600")
                    ui.button("后退", on_click=lambda: set_movement_state(2)).classes("bg-red-500 hover:bg-red-600")
                    ui.button("左移", on_click=lambda: set_movement_state(3)).classes("bg-blue-500 hover:bg-blue-600")
                    ui.button("右移", on_click=lambda: set_movement_state(4)).classes("bg-blue-500 hover:bg-blue-600")
                    ui.button("停止", on_click=lambda: set_movement_state(0)).classes("bg-gray-500 hover:bg-gray-600")

            # 右边：参数表
            with ui.card().classes("w-full h-full overflow-auto"):
                ui.label("LowCmd 参数表").classes("text-xl")
                lowcmd_table = ui.table(
                    columns=[
                        {'name': 'param', 'label': '参数', 'field': 'param', 'align': 'left'},
                        {'name': 'value', 'label': '值', 'field': 'value', 'align': 'right'}
                    ],
                    rows=[],
                    row_key='param'
                ).classes("w-full")

        # 第三行：3D场景 (全宽)
        with ui.card().classes("w-full h-[45vh] overflow-hidden"):
            ui.label("机器人3D空间").classes("text-xl")
            ui.add_body_html('''
                <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/build/three.min.js"></script>
                <script src="https://cdn.jsdelivr.net/npm/three@0.132.2/examples/js/controls/OrbitControls.min.js"></script>
                <script src="https://cdn.jsdelivr.net/npm/gsap@3.11.4/dist/gsap.min.js"></script>
            ''')
            scene_container = ui.html('''
                <div id="threejs-container" style="width:100%; height:100%; position:relative;">
                    <canvas id="threejs-canvas" style="display:block; width:100%; height:100%"></canvas>
                </div>
            ''').classes("w-full h-full")
            ui.run_javascript('''
                // 初始化Three.js场景
                const container = document.getElementById('threejs-container');
                const canvas = document.getElementById('threejs-canvas');
                canvas.width = container.clientWidth;
                canvas.height = container.clientHeight;

                const scene = new THREE.Scene();
                scene.background = new THREE.Color(0xf0f0f0);

                const camera = new THREE.PerspectiveCamera(75, container.clientWidth / container.clientHeight, 0.1, 1000);
                camera.position.set(5, 5, 5);
                camera.lookAt(0, 0, 0);

                const renderer = new THREE.WebGLRenderer({ canvas, antialias: true });
                renderer.setSize(canvas.width, canvas.height);

                // 添加辅助工具
                scene.add(new THREE.AxesHelper(5));
                scene.add(new THREE.GridHelper(10, 10));

                // 灯光
                scene.add(new THREE.AmbientLight(0x404040));
                const directionalLight = new THREE.DirectionalLight(0xffffff, 0.5);
                directionalLight.position.set(1, 1, 1);
                scene.add(directionalLight);

                // 创建小狗模型
                const dog = new THREE.Group();
                scene.add(dog);

                // 身体
                const body = new THREE.Mesh(
                    new THREE.CylinderGeometry(0.5, 0.5, 1.5, 16).rotateX(Math.PI/2),
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                dog.add(body);

                // 头部
                const head = new THREE.Mesh(
                    new THREE.SphereGeometry(0.4, 16, 16),
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                head.position.set(0.8, 0.2, 0);
                dog.add(head);

                // 鼻子
                const nose = new THREE.Mesh(
                    new THREE.SphereGeometry(0.1, 16, 16),
                    new THREE.MeshPhongMaterial({ color: 0x000000 })
                );
                nose.position.set(1.1, 0.2, 0);
                dog.add(nose);

                // 耳朵
                const earGeometry = new THREE.ConeGeometry(0.2, 0.4, 16);
                const leftEar = new THREE.Mesh(
                    earGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                leftEar.position.set(0.6, 0.5, 0.3);
                leftEar.rotation.z = -Math.PI/4;
                dog.add(leftEar);

                const rightEar = new THREE.Mesh(
                    earGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                rightEar.position.set(0.6, 0.5, -0.3);
                rightEar.rotation.z = -Math.PI/4;
                dog.add(rightEar);

                // 眼睛
                const eyeGeometry = new THREE.SphereGeometry(0.05, 16, 16);
                const leftEye = new THREE.Mesh(
                    eyeGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x000000 })
                );
                leftEye.position.set(0.9, 0.3, 0.15);
                dog.add(leftEye);

                const rightEye = new THREE.Mesh(
                    eyeGeometry,
                    new THREE.MeshPhongMaterial({ color: 0x000000 })
                );
                rightEye.position.set(0.9, 0.3, -0.15);
                dog.add(rightEye);

                // 腿
                const legGeometry = new THREE.CylinderGeometry(0.1, 0.1, 0.5, 8);
                const legPositions = [
                    { x: 0.3, y: -0.5, z: 0.3 }, { x: -0.3, y: -0.5, z: 0.3 },
                    { x: 0.3, y: -0.5, z: -0.3 }, { x: -0.3, y: -0.5, z: -0.3 }
                ];

                const legs = [];
                legPositions.forEach(pos => {
                    const leg = new THREE.Mesh(
                        legGeometry,
                        new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                    );
                    leg.position.set(pos.x, pos.y, pos.z);
                    dog.add(leg);
                    legs.push(leg);
                });

                // 尾巴
                const tail = new THREE.Mesh(
                    new THREE.CylinderGeometry(0.05, 0.05, 0.4, 8),
                    new THREE.MeshPhongMaterial({ color: 0x8B4513 })
                );
                tail.position.set(-0.8, -0.2, 0);
                tail.rotation.z = Math.PI/4;
                dog.add(tail);

                // 控制器和动画循环
                const controls = new THREE.OrbitControls(camera, renderer.domElement);
                controls.enableDamping = true;

                function animate() {
                    requestAnimationFrame(animate);
                    controls.update();
                    renderer.render(scene, camera);
                }
                animate();

                // 存储到全局
                window.robotScene = { scene, dog, legs, camera, renderer };

                // 响应式调整
                window.addEventListener('resize', () => {
                    camera.aspect = container.clientWidth / container.clientHeight;
                    camera.updateProjectionMatrix();
                    renderer.setSize(container.clientWidth, container.clientHeight);
                });
            ''')

    # 定时更新
    ui.timer(0.05, update_robot_position)  # 更快的更新频率使运动更平滑
    ui.timer(0.5, update_lowcmd_table)
    await connect_video_stream()


def set_movement_state(state):
    """设置机器人运动状态"""
    robot_status['movement_state'] = state
    if state == 0:
        send_gait_command(movement=0)
    else:
        send_gait_command(movement=state, velocity=0.5)


def update_robot_position():
    """更新机器人位置，实现平滑移动"""
    pos = robot_status['position']
    target = robot_status['target_position']
    movement_state = robot_status['movement_state']
    joints = robot_status['joints']

    # 根据运动状态更新目标位置
    speed = 0.05  # 移动速度
    if movement_state == 1:  # 前进
        target['y'] += speed
    elif movement_state == 2:  # 后退
        target['y'] -= speed
    elif movement_state == 3:  # 左移
        target['x'] -= speed
    elif movement_state == 4:  # 右移
        target['x'] += speed

    # 使用线性插值实现平滑移动
    pos['x'] += (target['x'] - pos['x']) * 0.1
    pos['y'] += (target['y'] - pos['y']) * 0.1

    status_label.text = f"状态: 位置({pos['x']:.2f}, {pos['y']:.2f}) 运动:{['停止', '前进', '后退', '左移', '右移'][movement_state]}"

    # 更新3D场景
    ui.run_javascript(f'''
        if (window.robotScene) {{
            const dog = window.robotScene.dog;
            const legs = window.robotScene.legs;

            // 使用GSAP实现平滑动画
            gsap.to(dog.position, {{
                x: {pos['x']},
                z: {pos['y']},
                duration: 0.2,
                ease: "power2.out"
            }});

            // 根据运动状态调整小狗朝向
            if ({movement_state} > 0) {{
                const targetRotation = {{
                    1: Math.PI,    // 前进
                    2: 0,          // 后退
                    3: Math.PI/2,  // 左移
                    4: -Math.PI/2  // 右移
                }}[{movement_state}];

                gsap.to(dog.rotation, {{
                    y: targetRotation,
                    duration: 0.3,
                    ease: "power2.out"
                }});
            }}

            // 模拟腿部动画
            const time = Date.now() * 0.001;
            if ({movement_state} > 0) {{
                legs.forEach((leg, i) => {{
                    const legAngle = Math.sin(time * 5 + i * Math.PI/2) * 0.2;
                    leg.rotation.z = legAngle;
                }});
            }} else {{
                legs.forEach(leg => {{
                    leg.rotation.z = 0;
                }});
            }}
        }}
    ''')


def update_lowcmd_table():
    if 'lowcmd' in robot_status:
        lowcmd = robot_status['lowcmd']
        rows = [
            {'param': '处理模式', 'value': '内置算法' if lowcmd['mode_pr'] == 0 else '用户处理'},
            {'param': '机器人类型', 'value': '机型4' if lowcmd['mode_machine'] == 4 else '机型6'},
            {'param': 'CRC校验', 'value': str(lowcmd['crc'])}
        ]

        # 添加电机命令数据（示例显示前3个电机）
        for i in range(3):
            motor = lowcmd['motor_cmd'][i]
            rows.extend([
                {'param': f'电机{i}模式', 'value': '刹车' if motor['mode'] == 0 else 'FOC'},
                {'param': f'电机{i}角度(rad)', 'value': f"{motor['q']:.3f}"},
                {'param': f'电机{i}速度(rad/s)', 'value': f"{motor['dq']:.3f}"},
                {'param': f'电机{i}力矩', 'value': f"{motor['tau']:.3f}"},
                {'param': f'电机{i}KP', 'value': f"{motor['kp']:.3f}"},
                {'param': f'电机{i}KD', 'value': f"{motor['kd']:.3f}"}
            ])

        lowcmd_table.rows = rows
        lowcmd_table.update()


async def connect_video_stream():
    global video_ws
    try:
        video_ws = await ui.run_javascript(f"""
            const socket = new WebSocket('{VIDEO_WS_URL}');
            socket.onopen = () => setInterval(() => socket.send('ping'), 15000);
            return socket;
        """)
    except Exception as e:
        print(f"连接视频流失败: {e}")


def send_gait_command(velocity=1.0, nominal_height=0.25, stance_duration=0.2,
                      stance_depth=0.0, swing_height=0.07, gait_type=0, movement=0):
    try:
        node = rclpy.create_node('gait_command_publisher')
        pub = node.create_publisher(Float32MultiArray, GAIT_MSG_TOPIC, 10)

        gait_msg = Float32MultiArray()
        gait_msg.data = [
            float(gait_type),
            float(velocity),
            float(nominal_height),
            float(stance_duration),
            float(stance_depth),
            float(swing_height),
            float(movement)
        ]

        pub.publish(gait_msg)
        node.destroy_node()

        robot_status['gait_params'] = {
            'gait_type': gait_type,
            'velocity': velocity,
            'nominal_height': nominal_height,
            'stance_duration': stance_duration,
            'stance_depth': stance_depth,
            'swing_height': swing_height,
            'movement': movement
        }

        ui.notify(f"已发送运动指令: {'停止' if movement == 0 else ['前进', '后退', '左移', '右移'][movement - 1]}")
    except Exception as e:
        ui.notify(f"ROS2 Error: {str(e)}", type='negative')


# 启动ROS2线程
threading.Thread(target=ros_thread, daemon=True).start()

# 启动NiceGUI应用
ui.run(
    title="Unitree Go2 Control Panel",
    port=5000,
    storage_secret="111",
    dark=False,
    host='0.0.0.0'
)